"""
历史数据处理模块
"""
from typing import Dict, List, Optional
import pandas as pd
import ccxt
from datetime import datetime

from .base import DataHandler
from utils.retry import retry_on_exception

class HistoricalDataHandler(DataHandler):
    """历史数据处理器"""
    
    def __init__(
        self,
        exchanges: Dict[str, ccxt.Exchange],
        symbols: List[str],
        timeframe: str,
        start_date: datetime,
        end_date: datetime
    ):
        """初始化历史数据处理器
        
        Args:
            exchanges: 交易所实例字典，键为交易所名称
            symbols: 交易对列表
            timeframe: 时间周期
            start_date: 开始日期
            end_date: 结束日期
        """
        super().__init__()
        
        self.exchanges = exchanges
        self.symbols = symbols
        self.timeframe = timeframe
        self.start_date = start_date
        self.end_date = end_date
        
        # 加载历史数据
        self._load_data()
    
    @retry_on_exception()
    def _load_data(self) -> None:
        """加载历史数据"""
        print("\n加载历史数据...")
        
        for exchange_name, exchange in self.exchanges.items():
            print(f"处理交易所: {exchange_name}")
            
            # 确保交易所键存在
            if exchange_name not in self._data:
                self._data[exchange_name] = {}
                self._latest_data[exchange_name] = {}
                self._current_idx[exchange_name] = {}
            
            for symbol in self.symbols:
                try:
                    print(f"正在从{exchange_name}获取{symbol}的历史数据...")
                    
                    # 获取历史数据
                    ohlcv = exchange.fetch_ohlcv(
                        symbol=symbol,
                        timeframe=self.timeframe,
                        since=int(self.start_date.timestamp() * 1000),
                        limit=None
                    )
                    
                    # 转换为DataFrame
                    df = pd.DataFrame(
                        ohlcv,
                        columns=['timestamp', 'open', 'high', 'low', 'close', 'volume']
                    )
                    
                    # 转换时间戳
                    df['timestamp'] = pd.to_datetime(df['timestamp'], unit='ms')
                    df.set_index('timestamp', inplace=True)
                    
                    # 过滤时间范围
                    df = df[
                        (df.index >= pd.Timestamp(self.start_date)) &
                        (df.index <= pd.Timestamp(self.end_date))
                    ]
                    
                    print(f"获取到{len(df)}条{symbol}的历史数据")
                    
                    # 存储数据
                    self._data[exchange_name][symbol] = df
                    self._latest_data[exchange_name][symbol] = []
                    self._current_idx[exchange_name][symbol] = 0
                    
                except Exception as e:
                    print(f"加载{exchange_name}的{symbol}数据时出错: {str(e)}")
                    raise
    
    def update_bars(self):
        """更新最新的K线数据"""
        for exchange_name in self.exchanges:
            for symbol in self.symbols:
                try:
                    if exchange_name not in self._data or symbol not in self._data[exchange_name]:
                        continue
                    
                    df = self._data[exchange_name][symbol]
                    current_idx = self._current_idx[exchange_name][symbol]
                    
                    if current_idx >= len(df):
                        self.continue_backtest = False
                        return
                    
                    latest_bar = df.iloc[current_idx].to_dict()
                    latest_bar['timestamp'] = df.index[current_idx]
                    self._latest_data[exchange_name][symbol].append(latest_bar)
                    self._current_idx[exchange_name][symbol] += 1
                    
                except Exception as e:
                    print(f"更新{exchange_name}的{symbol}数据时出错: {str(e)}")
                    raise
    
    def get_latest_bar(self, exchange: str, symbol: str) -> Optional[Dict]:
        """获取最新的一条数据"""
        try:
            latest_data = self._latest_data[exchange][symbol]
            return latest_data[-1] if latest_data else None
        except (KeyError, IndexError):
            return None
    
    def get_latest_bars(self, exchange: str, symbol: str, *, n: int = 50) -> List[Dict]:
        """获取最新的n条数据"""
        try:
            latest_data = self._latest_data[exchange][symbol]
            return latest_data[-n:] if len(latest_data) >= n else latest_data
        except (KeyError, IndexError):
            return []
